package set

import (
	"constraints"
	"encoding/json"
	"sync"
)

// arrSet 使用 slice 实现的集合
// 数据元素是可以排序的
type arrSet[T constraints.Ordered] struct {
	slice    []T
	comparer func(a, b T) bool
	mutex    *sync.RWMutex
	safe     bool
}

// NewArrSet 创建一个基于数组的集合
// 如果 comparer 不是 nil 则根据 comparer 进行排序，或者按照先后顺序进行排序
func NewArrSet[T constraints.Ordered](initCap int, comparer func(a, b T) bool, safe bool) *arrSet[T] {
	aset := &arrSet[T]{
		slice:    make([]T, 0, initCap),
		comparer: comparer,
		safe:     safe,
	}
	if safe {
		aset.mutex = new(sync.RWMutex)
	}
	return aset
}

// insertOne 按指定的顺序插入元素
// 此方法默认 element 是不存在的元素
func (aset *arrSet[T]) insertOne(ele T) {
	if aset.comparer == nil {
		aset.slice = append(aset.slice, ele)
		return
	}

	idx := -1
	for i := range aset.slice {
		if !aset.comparer(aset.slice[i], ele) {
			idx = i
			break
		}
	}

	if idx == -1 {
		aset.slice = append(aset.slice, ele)
	} else {
		// 通过 append 的链式操作插入元素
		aset.slice = append(aset.slice[:idx], append([]T{ele}, aset.slice[idx:]...)...)
	}
}

// Add 添加元素
func (aset *arrSet[T]) Add(elements ...T) {
	aset.maybeLock()
	defer aset.maybeUnlock()

	for i := range elements {
		if !aset.unsafeContains(elements[i]) {
			aset.insertOne(elements[i])
		}
	}
}

// Remove 移除元素
func (aset *arrSet[T]) Remove(elements ...T) {
	aset.maybeLock()
	defer aset.maybeUnlock()

	for i := range elements {
		idx := -1
		for j := range aset.slice {
			if aset.slice[j] == elements[i] {
				idx = j
				break
			}
		}
		if idx != -1 {
			aset.slice = append(aset.slice[:idx], aset.slice[idx+1:]...)
		}
	}
}

// unsafeContains 包含 ele 元素返回true，或者返回false
func (aset *arrSet[T]) unsafeContains(ele T) bool {
	for i := range aset.slice {
		if aset.slice[i] == ele {
			return true
		}
	}
	return false
}
func (aset *arrSet[T]) Contains(ele T) bool {
	aset.maybeRLock()
	defer aset.maybeRUnlock()
	return aset.unsafeContains(ele)
}

// Iterator 迭代
func (aset *arrSet[T]) Iterator(fn func(ele T) bool) {
	aset.maybeRLock()
	defer aset.maybeRUnlock()
	for i := range aset.slice {
		if !fn(aset.slice[i]) {
			return
		}
	}
}

// Len 元素数量
func (aset *arrSet[T]) Len() int {
	aset.maybeRLock()
	defer aset.maybeRUnlock()
	return len(aset.slice)
}

// Reset 清空容器
func (aset *arrSet[T]) Reset() {
	aset.maybeLock()
	defer aset.maybeUnlock()
	aset.slice = aset.slice[0:0]
}

// Slice 返回slice
func (aset *arrSet[T]) Slice() []T {
	aset.maybeRLock()
	defer aset.maybeRUnlock()
	newSlice := make([]T, aset.Len())
	copy(newSlice, aset.slice)
	return newSlice
}

// Clone 克隆
func (aset *arrSet[T]) Clone() *arrSet[T] {
	aset.maybeRLock()
	defer aset.maybeRUnlock()

	newSet := NewArrSet(cap(aset.slice), aset.comparer, aset.safe)
	copy(newSet.slice, aset.slice)
	return newSet
}

// Union 并集
func (aset *arrSet[T]) Union(set *arrSet[T]) *arrSet[T] {
	newSet := aset.Clone()
	newSet.Add(set.slice...)
	return newSet
}

// Intersect 交集
func (aset *arrSet[T]) Intersect(set *arrSet[T]) *arrSet[T] {
	newSet := NewArrSet(cap(aset.slice), aset.comparer, aset.safe)
	aset.Iterator(func(ele T) bool {
		if set.Contains(ele) {
			newSet.Add(ele)
		}
		return true
	})

	return newSet
}

// Difference 差集
func (aset *arrSet[T]) Difference(set *arrSet[T]) *arrSet[T] {
	newSet := NewArrSet(cap(aset.slice), aset.comparer, aset.safe)
	aset.Iterator(func(ele T) bool {
		if !set.Contains(ele) {
			newSet.Add(ele)
		}
		return true
	})

	return newSet
}

// SymmetricDifference 对称差集
func (aset *arrSet[T]) SymmetricDifference(set *arrSet[T]) *arrSet[T] {
	newSet1 := aset.Difference(set)
	newSet2 := set.Difference(aset)
	return newSet1.Union(newSet2)
}

// MarshalText
func (aset *arrSet[T]) MarshalText() ([]byte, error) {
	return json.Marshal(aset.slice)
}

// UnmarshalText
func (aset *arrSet[T]) UnmarshalText(data []byte) error {
	elements := make([]T, 0, 1024)
	if err := json.Unmarshal(data, &elements); err != nil {
		return err
	}
	aset.Add(elements...)
	return nil
}

// UnmarshalJSON
func (aset *arrSet[T]) UnmarshalJSON(data []byte) error {
	return aset.UnmarshalText(data)
}

func (hset *arrSet[T]) maybeLock() {
	if hset.safe {
		hset.mutex.Lock()
	}
}

func (hset *arrSet[T]) maybeUnlock() {
	if hset.safe {
		hset.mutex.Unlock()
	}
}

func (hset *arrSet[T]) maybeRLock() {
	if hset.safe {
		hset.mutex.RLock()
	}
}

func (hset *arrSet[T]) maybeRUnlock() {
	if hset.safe {
		hset.mutex.RUnlock()
	}
}
